import torch
from torch.utils.data import DataLoader
from src.models.modnet import MODNet
from src.trainer import supervised_training_iter
from dataset.data import SegmentPPAlignImage
import dataset.joint_transforms as joint_transforms

bs = 16         # batch size
lr = 0.01       # learn rate
epochs = 40     # total epochs

modnet = torch.nn.DataParallel(MODNet()).cuda()
optimizer = torch.optim.SGD(modnet.parameters(), lr=lr, momentum=0.9)
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=int(0.25 * epochs), gamma=0.1)

joint_transform = joint_transforms.Compose([
    joint_transforms.FreeScale((672, 576)),
    joint_transforms.RandomHorizontallyFlip(),
    joint_transforms.RandomRotate(2),
    joint_transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.05),
    joint_transforms.RandomSizedCrop((672, 576), (0.95, 1.05), 30),
])
dataset = SegmentPPAlignImage(root='/home/paramai/data/removebg-align/combine_train',
                            joint_transform=joint_transform, unknown_range=(0, 254))
dataloader = DataLoader(dataset, bs, lr=lr, num_workers=4, drop_last=True)     # NOTE: please finish this function

for epoch in range(0, epochs):
    for idx, (image, trimap, gt_matte) in enumerate(dataloader):
        semantic_loss, detail_loss, matte_loss = \
            supervised_training_iter(modnet, optimizer, image, trimap, gt_matte)
    lr_scheduler.step()